from transformers import AutoModelForCausalLM, AutoTokenizer
import json
import requests


# Function Call机制：通过解析来自模型输出中的特殊标记（如function_call:），触发对应的功能调用。

# 加载预训练模型和分词器
model_name = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def generate_response(input_text, history=[]):
    # 根据历史对话构建上下文
    context = "\n".join(history) + "\n" if history else ""
    inputs = tokenizer.encode(context + input_text + tokenizer.eos_token, return_tensors='pt')
    
    outputs = model.generate(inputs, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
    return response

def function_call_dispatcher(function_call):
    """根据函数调用请求执行相应的操作"""
    try:
        # 解析函数调用请求
        function_name = function_call['name']
        arguments = json.loads(function_call['arguments'])

        if function_name == "get_weather":
            response = requests.get("http://api.weatherapi.com/v1/current.json", params={"key": "your_api_key", "q": arguments["location"]})
            weather_data = response.json()
            return f"当前{arguments['location']}的温度是{weather_data['current']['temp_c']}度 Celsius。"
        elif function_name == "search_wikipedia":
            # 这里简化处理，实际应用中需要更复杂的逻辑
            return "Wikipedia搜索结果（此处省略）。"
        else:
            return "未知的功能调用。"
    except Exception as e:
        return str(e)

def smart_agent(input_text, history=[]):
    # 生成响应
    response = generate_response(input_text, history)
    print(f"初步响应: {response}")
    
    # 检查是否包含函数调用请求
    if 'function_call' in response:
        function_call = json.loads(response.split('function_call: ')[-1])
        execution_result = function_call_dispatcher(function_call)
        response += f"\n执行结果: {execution_result}"
    
    # 更新历史记录
    history.append(f"用户: {input_text}")
    history.append(f"助手: {response}")
    return response

if __name__ == "__main__":
    print("开始与智能助手对话（输入'exit'退出）：")
    history = []
    while True:
        input_text = input("你: ")
        if input_text.lower() == 'exit':
            break
        response = smart_agent(input_text, history)
        print(f"AI: {response}")